import numpy as np
import einops.layers.torch as elt

#载入数据
x_train = np.load("/home/wangwei83/Desktop/pytorch-dl-cv/dataset/mnist/x_train.npy")
y_train_label = np.load("/home/wangwei83/Desktop/pytorch-dl-cv/dataset/mnist/y_train_label.npy")

x_train = np.expand_dims(x_train,axis=1)	#在指定维度上进行扩充
print(x_train.shape)

import torch
import torch.nn as nn
import numpy as np
import einops.layers.torch as elt


class MnistNetword(nn.Module):
    def __init__(self):
        super(MnistNetword, self).__init__()
        # 前置的特征提取模块


        self.convs_stack = nn.Sequential(
            nn.Conv2d(1, 12, kernel_size=7),  # 第一个卷积层
            nn.ReLU(),
            nn.Conv2d(12, 24, kernel_size=5),  # 第二个卷积层
            nn.ReLU(),
            nn.Conv2d(24, 6, kernel_size=3)  # 第三个卷积层
        )
        # 最终分类器层
        self.logits_layer = nn.Linear(in_features=1536, out_features=10)


    def forward(self, inputs):
        image = inputs
        x = self.convs_stack(image)

        # elt.Rearrange的作用是对输入数据维度进行调整，读者可以使用torch.nn.Flatten函数完成此工作
        x = elt.Rearrange("b c h w -> b (c h w)")(x)
        logits = self.logits_layer(x)
        return logits


model = MnistNetword()
torch.save(model, "model.pth")
